import numpy as np
import matplotlib.pyplot as plt
from enthought.mayavi import mlab


def minmax(x, y, z):
    import lxc
    flat = [j for i in z.tolist() for j in i]
    zmin, zmax = lxc.minmax(flat) 
    return [x[0][0],x[-1][-1],y[0][0],y[-1][-1], zmin, zmax]
    
def draw3D(x, y, z):
    #x, y = np.ogrid[-2:2:20j, -2:2:20j]
    #z = x * np.exp(-x ** 2 - y ** 2)

    pl = mlab.surf(x, y, z, warp_scale="auto")
    mlab.axes(xlabel='x', ylabel='y', zlabel='z', nb_labels=5,
              ranges=minmax(x, y, z))
    mlab.outline(pl)

    mlab.show()

def draw3DContour(x, y, z, contourlines=28):
    #pl = mlab.surf(x, y, z, warp_scale="auto")
    pl = mlab.contour_surf(x, y, z, contours=contourlines, warp_scale=10)
    mlab.axes(xlabel='x', ylabel='y', zlabel='z', nb_labels=5,
              ranges=minmax(x, y, z))
    mlab.outline(pl)
    mlab.show()

def draw2DContour(x, y, z):
    #fig = plt.figure()
    mlab.imshow(z, line_width=3.0)
    mlab.axes(xlabel='x', ylabel='y', zlabel='z', nb_labels=5,
              ranges=minmax(x, y, z))
    mlab.show()

def example(degree = 0.01):
    #x, y = np.ogrid[-2:2:20j, -2:2:20j]
    x, y = np.mgrid[-2:2+degree:degree, -2:2+degree:degree]
    z = x * np.exp(-x ** 2 - y ** 2)
    return x, y, z

if __name__ == "__main__":
    k = example()
    draw3D(*k)
    #draw3DContour(x, y, z)
    #draw2DContour(*k)
